import csv
import numpy
import random
import time


def load(filename, trainingSize=0.75, rseed=-1, normalize=True, transpose=True):
    '''
    Load data from tab delimited file.  Then, return normalized, transposed
    training and test sets.
    :param filename: file to load data from
    :type filename: str
    :param trainingSize: percent of data to use for training
    :type trainingSize: float
    :param rseed: seed to random number generator
    :type rseed: int
    :return: training set and test set
    :rtype: numpy.array tuple
    '''

    # read tab delimited data from file
    reader = csv.reader(open(filename,'rb'), delimiter='\t')
    # cast data to list of lists
    dataAsList = list(reader)
    # convert to matrix
    matWithLabels= numpy.array(dataAsList)
    
    print('Original matrix shape: ', matWithLabels.shape)
    # find first row that contains "control" data
    controlIndex = 0
    for i in range(2,matWithLabels.shape[0]):
        if matWithLabels[i,1]=='Control':
            controlIndex = i
            break

    # trim labels and cast to float matrix
    dataMat = numpy.array(matWithLabels[2:,4:]).astype('float')
    print('Original Data shape: ', dataMat.shape)
    
    # extract case and control matrices
    caseMat = dataMat[:controlIndex-2,:]
    controlMat = dataMat[controlIndex-2:,:]
    print('Case Data shape: ', caseMat.shape)
    print('Control Data shape: ', controlMat.shape)

    # discrete case / control matrices
    discreteCaseMat = numpy.array(matWithLabels[2:controlIndex,0:3])
    discreteControlMat = numpy.array(matWithLabels[controlIndex:,0:3])
    
    # seed random number generator
    if rseed<0:
        rseed = int(time.time()) # use time as seed
    random.seed(rseed)

    # create indices for columns in case and control matrices
    caseIndices = list(range(caseMat.shape[0]))
    controlIndices = list(range(controlMat.shape[0]))

    # randomly shuffle indices
    random.shuffle(caseIndices)
    random.shuffle(controlIndices)

    # find last indices for training sets based on 'trainingSize'
    endIndexCase = int(caseMat.shape[0]*trainingSize)
    endIndexControl = int(controlMat.shape[0]*trainingSize)

    # grab 'trainingSize' of each matrix for training set
    trainingCaseMat = caseMat[caseIndices[:endIndexCase],:]
    trainingControlMat = controlMat[controlIndices[:endIndexControl],:]   

    # grab '1-trainingSize' of each matrix for testing set
    testingCaseMat = caseMat[caseIndices[endIndexCase:],:]
    testingControlMat = controlMat[controlIndices[endIndexControl:],:]

    # concatenate matrices to construct training and test matrices
    trainingMat = numpy.concatenate((trainingCaseMat,trainingControlMat),axis=0)
    testingMat = numpy.concatenate((testingCaseMat,testingControlMat),axis=0)
    print('(\'Training Data shape: \', {}; n_case: {}; n_control: {})'.format(trainingMat.shape, trainingCaseMat.shape[0], trainingControlMat.shape[0]))
    print('(\'Test Data shape: \', {}; n_case: {}; n_control: {})'.format(testingMat.shape, testingCaseMat.shape[0], testingControlMat.shape[0]))

    #print('Case[:',endIndexCase, ']=', discreteMat[caseIndices[:endIndexCase],:])
    #print('Control[:',endIndexControl, ']=', discreteMat[controlIndices[:endIndexControl],:])

    #print('Case[:',endIndexCase, ']=', discreteMat[caseIndices[endIndexCase:],:])
    #print('Control[:',endIndexControl, ']=', discreteMat[controlIndices[endIndexControl:],:])

    #grab the Age, Case/Control and M/F discrete values for training and testing matrices    
    trainingDiscreteMat = numpy.concatenate((discreteCaseMat[caseIndices[:endIndexCase],:], discreteControlMat[controlIndices[:endIndexControl],:]),axis=0)
    testingDiscreteMat = numpy.concatenate((discreteCaseMat[caseIndices[endIndexCase:],:], discreteControlMat[controlIndices[endIndexControl:],:]),axis=0)

    # normalize matrices
    if(normalize):
        print('Normalizing matrices')
        normalizer = trainingMat.mean()
        trainingMat = trainingMat / normalizer
        testingMat = testingMat / normalizer

    # ALT-Normalize: (appears worse than divide by element-wise mean)
    #     Idea was to normalize each protein's level based on community levels of said protein
    #print 'Normalizing all protein levels to ratio [0,1]'
    # NOTE: After Transpose, each row of matrix is a given protein with 234 column-patients
    #for x in range(trainingMat.shape[0]):
    #    trainingMat[x] = trainingMat[x] / trainingMat[x].max()
    #for x in range(testingMat.shape[0]):
    #    testingMat[x] = testingMat[x] / testingMat[x].max()


    # transpose matrix
    if(transpose):
        print('Transposing matrices (patients are now column vectors)')
        trainingMat = trainingMat.transpose()
        testingMat = testingMat.transpose()
        trainingDiscreteMat = trainingDiscreteMat.transpose()
        testingDiscreteMat = testingDiscreteMat.transpose()
        print('Transpose trainingMat.shape: ', trainingMat.shape, 'testingMat.shape: ', testingMat.shape)
        print('Transpose trainingDiscreteMat.shape: ', trainingDiscreteMat.shape, 'testingDiscreteMat.shape: ', testingDiscreteMat.shape)

    return (trainingMat, testingMat, trainingDiscreteMat, testingDiscreteMat)


